from tests.pytests.helpers.codegen_helpers import compile_chunk, assert_code_contains


def test_range_loop():
    def sum_range(n: int) -> int:
        s = 0
        for i in range(n):
            s += i
        return s

    block = compile_chunk(sum_range)
    assert_code_contains(
        block,
        """
        int sum_range(int n) {
            int s = 0;
            for (int i = 0; i < n; ++i) {
                s += i;
            }
            return s;
        }
        """,
    )


def test_while_loop():
    def sum_while(n: int) -> float:
        s = 0.5
        i = 1
        while i < n:
            s += float(i)
            i += 1
        return s

    block = compile_chunk(sum_while)
    assert_code_contains(
        block,
        """
        double sum_while(int n) {
            double s = 0.5;
            int    i = 1;
            while (i < n) {
                s += static_cast<double>(i);
                i += 1;
            }
            return s;
        }
        """,
    )


def test_for_each_over_list():
    def sum_list(xs: list[int]) -> int:
        s = 0
        for x in xs:
            s += x
        return s

    block = compile_chunk(sum_list)
    assert_code_contains(
        block,
        """
        int sum_list(const mc::array<int>& xs) {
            int s = 0;
            for (const auto& x : xs) {
                s += x;
            }
            return s;
        }
        """,
    )


def test_range_loop_with_correct_type():
    """测试 range 循环使用正确的类型而不是 auto"""

    def sum_range_typed(n: int) -> int:
        total = 0
        for i in range(n):
            doubled = i * 2  # i 应该被推断为 int 类型
            total += doubled
        return total

    block = compile_chunk(sum_range_typed)
    assert_code_contains(
        block,
        """
        int sum_range_typed(int n) {
            int total = 0;
            for (int i = 0; i < n; ++i) {
                int doubled = i * 2;
                total += doubled;
            }
            return total;
        }
        """,
    )


def test_range_loop_with_start_and_step():
    """测试带起始值和步长的 range 循环"""

    def sum_range_step(start: int, end: int, step: int) -> int:
        total = 0
        for i in range(start, end, step):
            total += i
        return total

    block = compile_chunk(sum_range_step)
    assert_code_contains(
        block,
        """
        int sum_range_step(int start, int end, int step) {
            int total = 0;
            for (int i = start; i < end; i += step) {
                total += i;
            }
            return total;
        }
        """,
    )


def test_for_each_over_string():
    """测试字符串迭代"""

    def count_chars(text: str) -> int:
        count = 0
        for char in text:
            if char == "a":
                count += 1
        return count

    block = compile_chunk(count_chars)
    assert_code_contains(
        block,
        """
        int count_chars(const std::string& text) {
            int count = 0;
            for (const auto& char : text) {
                if (char == "a") {
                    count += 1;
                }
            }
            return count;
        }
        """,
    )


def test_for_each_over_dict():
    """测试字典迭代（迭代键）"""

    def sum_dict_keys(data: dict[str, int]) -> int:
        total = 0
        for key in data:
            total += len(key)
        return total

    block = compile_chunk(sum_dict_keys)
    assert_code_contains(
        block,
        """
        int sum_dict_keys(const mc::dict& data) {
            int total = 0;
            for (const auto& key : mcpy::dict_keys<std::string>(data)) {
                total += key.size();
            }
            return total;
        }
        """,
    )


def test_for_each_tuple_unpacking():
    """测试元组解包循环"""

    def sum_pairs(pairs: list[tuple[int, int]]) -> int:
        total = 0
        for x, y in pairs:
            total += x + y
        return total

    block = compile_chunk(sum_pairs)
    expected = """
    int sum_pairs(const mc::array<std::tuple<int, int>>& pairs) {
        int total = 0;
        for (const auto& [x, y] : pairs) {
            total += x + y;
        }
        return total;
    }
    """
    assert_code_contains(block, expected)


def test_for_each_multi_variable_unpacking():
    """测试多变量解包循环"""

    def sum_triples(triples: list[tuple[int, int, int]]) -> int:
        total = 0
        for x, y, z in triples:
            total += x + y + z
        return total

    block = compile_chunk(sum_triples)
    expected = """
    int sum_triples(const mc::array<std::tuple<int, int, int>>& triples) {
        int total = 0;
        for (const auto& [x, y, z] : triples) {
            total += x + y + z;
        }
        return total;
    }
    """
    assert_code_contains(block, expected)


def test_nested_loops():
    """测试嵌套循环"""

    def nested_sum(n: int, m: int) -> int:
        total = 0
        for i in range(n):
            for j in range(m):
                total += i * j
        return total

    block = compile_chunk(nested_sum)
    assert_code_contains(
        block,
        """
        int nested_sum(int n, int m) {
            int total = 0;
            for (int i = 0; i < n; ++i) {
                for (int j = 0; j < m; ++j) {
                    total += i * j;
                }
            }
            return total;
        }
        """,
    )


def test_loop_with_variable_declarations():
    """测试循环中的变量声明"""

    def loop_with_vars(numbers: list[int]) -> int:
        result = 0
        for num in numbers:
            doubled = num * 2  # 应该推断为 int
            squared = doubled * doubled  # 应该推断为 int
            result += squared
        return result

    block = compile_chunk(loop_with_vars)
    assert_code_contains(
        block,
        """
        int loop_with_vars(const mc::array<int>& numbers) {
            int result = 0;
            for (const auto& num : numbers) {
                int doubled = num * 2;
                int squared = doubled * doubled;
                result += squared;
            }
            return result;
        }
        """,
    )


def test_float_range_loop():
    """测试浮点数 range 循环（虽然 Python 的 range 只支持整数，但测试类型推断）"""

    def float_operations() -> float:
        total = 0.0
        for i in range(5):
            float_val = float(i) * 1.5  # i 是 int，转换为 float
            total += float_val
        return total

    block = compile_chunk(float_operations)
    assert_code_contains(
        block,
        """
        double float_operations() {
            double total = 0.0;
            for (int i = 0; i < 5; ++i) {
                double float_val = static_cast<double>(i) * 1.5;
                total += float_val;
            }
            return total;
        }
        """,
    )
